from agentscope.message import Msg
from Server.managers.nodes.baseGraphNode import BaseGraphNode
from AlgoriAgent.projects.commonAgent.agents.commonAgent import CommonAgent
from .sessionInfo import SessionInfo
import json
import uuid
import eventlet
from eventlet import Timeout
class ToolExecuteNode(BaseGraphNode):
    def __init__(self, id, data, graphNodeManager, random_key):
        super().__init__(id, data, graphNodeManager, random_key)
        self.route = data['data']['route']
        self.timeout = data['data'].get('timeout',10)
        self.flow_next_nodes_failed=[]
        self.flow_next_nodes_success = []
        self.session_pending_events = {}  # {session_id: {correlation_id: event}}
        

    def add_flow_next(self, node_id, **kwargs):
        if 'successed' in kwargs['sourceHandle']:
            self.flow_next_nodes_success.append(node_id)
        elif 'failed' in kwargs['sourceHandle']:
            self.flow_next_nodes_failed.append(node_id)
        else:
            self.flow_next_nodes.append(node_id)
    def create_tool_execute_lisener(self, session_info: SessionInfo, actureCall: dict, timeout_seconds: int):
        # 确保该会话的事件字典存在
        if session_info.id not in self.session_pending_events:
            self.session_pending_events[session_info.id] = {}
            
        correlation_id = str(uuid.uuid4())
        event = eventlet.event.Event()
        # 将事件存储在会话特定的字典中
        self.session_pending_events[session_info.id][correlation_id] = event
        
        print(f"Setting up listener for session {session_info.id}, correlation_id: {correlation_id}")
        
        @session_info.ws.on(self.route, namespace=session_info.full_url)
        def handle_channel_response(msg):
            print(f"Received message for session {session_info.id}: {msg}")
            received_correlation_id = msg.get('correlation_id')
            if not received_correlation_id:
                print(f"No correlation_id in message")
                return
                
            # 检查该会话的事件字典中是否存在对应的事件
            session_events = self.session_pending_events.get(session_info.id, {})
            if received_correlation_id in session_events:
                print(f"Found matching event for session {session_info.id}, correlation_id {received_correlation_id}")
                pending_event = session_events[received_correlation_id]
                data = msg.get('data')
                print(f"Sending data to event: {data}")
                pending_event.send(data)
            else:
                print(f"No pending event found for session {session_info.id}, correlation_id {received_correlation_id}")
        
        try:
            print(f"Emitting message for session {session_info.id}, correlation_id: {correlation_id}")
            session_info.ws.emit(
                self.route, 
                {
                    'data': actureCall, 
                    'correlation_id': correlation_id,
                    'route': self.route
                }, 
                room=session_info.sid, 
                namespace=session_info.full_url
            )
            
            with Timeout(timeout_seconds):
                print(f"Waiting for response for session {session_info.id} with timeout {timeout_seconds}s")
                response_data = event.wait()
                print(f"Received response for session {session_info.id}: {response_data}")
                return response_data
                
        except Timeout:
            print(f"Timeout waiting for response for session {session_info.id}")
            return None
        except Exception as e:
            print(f"Unexpected error for session {session_info.id}: {e}")
            return None
        finally:
            # 清理会话特定的事件
            if session_info.id in self.session_pending_events:
                session_events = self.session_pending_events[session_info.id]
                if correlation_id in session_events:
                    del session_events[correlation_id]
                # 如果会话的事件字典为空，也可以清理会话
                if not session_events:
                    del self.session_pending_events[session_info.id]

    def __del__(self):
        # 清理所有会话的事件
        self.session_pending_events.clear()

    def execute(self, session_info = None):
        functionCall = None
        if 'function-call' in self.dict['in']:
            functionFrom = self.dict['in']['function-call']
            prenode =self.graphNodeManager.nodesByIdDict[functionFrom['nodeId']]
            functionCall = prenode.get_newest_output_in_session(session_info, functionFrom['handle'])
        toolName = None
        if 'tool' in self.dict['in']:
            toolFrom = self.dict['in']['tool']
            prenode =self.graphNodeManager.nodesByIdDict[toolFrom['nodeId']]
            if (prenode.type == 'tool'):
                toolName =prenode.get_tool_name()
            else: toolName=None
            # tool = prenode.get_newest_output_in_session(session_info, toolFrom['handle'])

        actureCall=None
        try:
            calls = []
            if(type(functionCall)==str):
                calls = json.loads(functionCall)
            else: calls = functionCall
            for call in calls:
                if call['name'] == toolName:
                    actureCall = call
                    break
        except Exception as e:
            print(e)
            actureCall = None

        print("execute Tool",actureCall)
        if session_info.id not in self.dict['out']: self.dict['out'][session_info.id] = {}

        if('result-msg' not in self.dict['out'][session_info.id]): self.dict['out'][session_info.id]['result-msg'] = []
        if(functionCall is None):
            self.dict['out'][session_info.id]['result-msg'].append("function call not found, check your project edge links")
        if (toolName is None):
            self.dict['out'][session_info.id]['result-msg'].append("tool not found, check your project edge links")
        if (actureCall is None):
            self.dict['out'][session_info.id]['result-msg'].append("tool name not match, this is normal when agent not pose a matched tool call")
        if (functionCall is None or toolName is None or actureCall is None):
            self.flow_next(session_info, no_execute=True)
            return

        try:
            response = self.create_tool_execute_lisener(session_info, actureCall, self.timeout)
        except Exception as e:
            response = None
            print(e)
            

        if response is not None:
            self.dict['out'][session_info.id]['result-msg'].append("Tool "+toolName+" executed successfully with result: "+ str(response))
            self.executed_num+=1
            self.flow_next(session_info, True)
            # 在这里处理 response
        else:
            print("Tool"+toolName+"执行失败")
            self.dict['out'][session_info.id]['result-msg'].append("Tool "+toolName+" execution failed")
            self.flow_next(session_info, False)
    

    def flow_next_process(self, session_info: SessionInfo, nodes:list[str]):
        for node in nodes:
            session_info.flow_prev_nodes_executed_number[node][self.id] += 1
            self.graphNodeManager.add_to_execution_queue(node, session_info)
            
    def get_all_flow_next_nodes(self):
        return self.flow_next_nodes + self.flow_next_nodes_success + self.flow_next_nodes_failed



    def flow_next(self, session_info: SessionInfo, success:bool = True, no_execute:bool = False):
        if no_execute:
            self.flow_next_process(session_info, self.flow_next_nodes)
        elif success:
            self.flow_next_process(session_info, self.flow_next_nodes_success)
        else:
            self.flow_next_process(session_info, self.flow_next_nodes_failed)
